from models.guided_ddpm_unet import UNetModel
from utils.galaxy_data_utils.transform_util import *
import torch
import numpy as np


class PpgUnetModel(UNetModel):

    def __init__(self, image_size, in_channels, *args, **kwargs):
        assert in_channels == 1, "mri image is considered"
        super().__init__(image_size, in_channels, *args, **kwargs)


    def forward(self, x, timesteps, psf):
        """
        前向传播方法，用于计算模型的输出。

        :param x: 形状为 [N x 2 x H x W] 的输入张量，代表时刻 t 的 x_t。
        :param timesteps: 一批时间步索引。
        :param image_dir: 形状为 [N x 2 x H x W] 的张量，代表脏图像（dirty img）。?
        :param psf: 点扩散函数图像
        :return: 噪声估计结果。
        """
        # print("ppgunetmodel-forward:")
        # print(x.shape)  # (1, 1, 256, 256)
        # print(timesteps.shape)  # 1
        # print(psf.shape)  # (1, 1, 256, 256)

        # 调用父类的 forward 方法，传入拼接后的输入、时间步和可见性数据
        output = super().forward(x, timesteps, psf)
        
        # 返回模型的输出，即噪声估计结果
        return output